#!/usr/bin/env python3
import cv2
import rclpy
from rclpy.node import Node
import cv_bridge
from sensor_msgs.msg import Image
from sensor_msgs.msg import CameraInfo
from sensor_msgs.msg import CompressedImage
import os
import numpy as np

camera_matrix = np.array(
    [
        [1572.569610, 0.000000, 607.879891],
        [0.000000, 1570.246204, 379.644093],
        [0.000000, 0.000000, 1.000000],
    ]
)
dist_coeffs = np.array([0.083966, -0.071645, 0.007297, 0.001435])



class CameraReader(Node):
    def __init__(self):
        super().__init__("camera_reader")
        self.declare_parameter("device", "undefine")
        self.declare_parameter("width", 1280)
        self.declare_parameter("height", 720)
        self.declare_parameter("fps", 60)
        self.declare_parameter("flip_horizontal", False)
        self.declare_parameter("flip_vertical", False)
        self.declare_parameter("channel_compress", False)
        device = self.get_parameter("device").get_parameter_value().string_value
        width = self.get_parameter("width").get_parameter_value().integer_value
        height = self.get_parameter("height").get_parameter_value().integer_value
        fps = self.get_parameter("fps").get_parameter_value().integer_value
        self.flip_horizontal = (
            self.get_parameter("flip_horizontal").get_parameter_value().bool_value
        )
        self.flip_vertical = (
            self.get_parameter("flip_vertical").get_parameter_value().bool_value
        )
        self.compress_channel_enable = (
            self.get_parameter("channel_compress").get_parameter_value().bool_value
        )

        if self.flip_horizontal and self.flip_vertical:
            self.flip = -1
        elif self.flip_horizontal:
            self.flip = 1
        elif self.flip_vertical:
            self.flip = 0
        else:
            self.flip = -2

        if device != "undefine":
            self.cap = cv2.VideoCapture(device, cv2.CAP_V4L)
        else:
            self.cap = cv2.VideoCapture(0, cv2.CAP_V4L)
        if not self.cap.isOpened():
            self.get_logger().error("Failed to open camera")
            exit(1)

        self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc("M", "J", "P", "G"))
        # self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc("Y", "U", "Y", "2"))
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
        self.cap.set(cv2.CAP_PROP_FPS, fps)
        self.create_timer(1.0 / fps, self.read_frame_callback)
        self.pub_image = self.create_publisher(Image, "camera_image", 1)
        if self.compress_channel_enable:
            self.pub_image_compress = self.create_publisher(
                CompressedImage, "camera_image_compress", 1
            )
        self.pub_info = self.create_publisher(CameraInfo, "camera_info", 1)
        self.bridge = cv_bridge.CvBridge()

    def read_frame_callback(self):
        frame = self.read_frame()
        if frame is None:
            self.get_logger().error("Failed to read frame")
            return
        self.pub_image.publish(
            self.bridge.cv2_to_imgmsg(frame, encoding="bgr8")
        )
        if self.compress_channel_enable:
            self.pub_image_compress.publish(
                self.bridge.cv2_to_compressed_imgmsg(frame, "jpg")
            )
        info = CameraInfo()
        info.header.stamp = self.get_clock().now().to_msg()
        info.header.frame_id = "camera"
        info.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        info.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        info.k = camera_matrix.reshape(-1).tolist()
        info.d = dist_coeffs.tolist()
        self.pub_info.publish(info)

    def read_frame(self):
        ret, frame = self.cap.read()
        if ret:
            if self.flip != -2:
                frame = cv2.flip(frame, self.flip)
            return frame
        else:
            return None

    def __del__(self):
        self.cap.release()


def main(args=None):
    rclpy.init(args=args)
    node = CameraReader()
    node.get_logger().info("Camera Reader node started")
    rclpy.spin(node)
    rclpy.shutdown()


if __name__ == "__main__":
    main()
